#!/usr/bin/env python
# encoding: utf-8

import torch
import argparse
import numpy as np


parser = argparse.ArgumentParser()
parser.add_argument('--data_list', help='data list', type=str, default="data/tmp_list")
parser.add_argument('--output_test_path', help='embedding dim', type=str, default="data/test_list")
parser.add_argument('--output_enroll_path', help='embedding dim', type=str, default="data/enroll_list")
args = parser.parse_args()

audio_seq, speaker = torch.load(args.data_list)
np.random.shuffle(audio_seq)

test_seq = audio_seq[:1000]
enroll_seq = audio_seq[1000:]

print("test length: ", len(test_seq))
print("enroll length: ", len(enroll_seq))

torch.save((test_seq, speaker), args.output_test_path)
torch.save((enroll_seq, speaker), args.output_enroll_path)
